"""
时间序列ARIMA预测模型

问题描述：
    ARIMA(Auto Regressive Integrated Moving Average)是一种经典的时间序列预测模型，
    适用于具有趋势和季节性的时间序列数据。ARIMA模型由三个参数(p,d,q)组成：
    - p: 自回归项数
    - d: 差分阶数  
    - q: 移动平均项数

算法原理：
    ARIMA模型结合了三个组件：
    1. AR(p) - 自回归：当前值依赖于前p个历史值
    2. I(d) - 积分：对时间序列进行d次差分以消除趋势
    3. MA(q) - 移动平均：当前值依赖于前q个预测误差
    
    数学表达式：
    (1-φ₁L-φ₂L²-...-φₚLᵖ)(1-L)ᵈXₜ = (1+θ₁L+θ₂L²+...+θqLᵠ)εₜ

核心概念：
    - 平稳性：时间序列的统计特性不随时间变化
    - 差分：消除时间序列的趋势性
    - 自相关函数(ACF)和偏自相关函数(PACF)：用于确定模型参数

应用场景：
    时间序列预测广泛应用于：
    - 股票价格预测
    - 销售额预测
    - 经济指标预测
    - 气温预测
    - 疫情传播预测

作者：斯黄
日期：2025年3月4日
"""

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.stattools import adfuller
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
import warnings
warnings.filterwarnings('ignore')

# 生成示例时间序列数据：模拟月销售数据
np.random.seed(42)
dates = pd.date_range('2020-01-01', periods=60, freq='M')
trend = np.linspace(100, 200, 60)  # 线性趋势
seasonal = 20 * np.sin(2 * np.pi * np.arange(60) / 12)  # 季节性
noise = np.random.normal(0, 10, 60)  # 随机噪声
sales = trend + seasonal + noise

# 创建时间序列DataFrame
ts_data = pd.DataFrame({'date': dates, 'sales': sales})
ts_data.set_index('date', inplace=True)

print("原始时间序列数据预览:")
print(ts_data.head())

# 绘制原始时间序列
plt.figure(figsize=(12, 8))

plt.subplot(2, 2, 1)
plt.plot(ts_data.index, ts_data['sales'], marker='o')
plt.title('原始时间序列数据')
plt.xlabel('时间')
plt.ylabel('销售额')
plt.grid(True)

# 平稳性检验
def check_stationarity(timeseries):
    result = adfuller(timeseries)
    print(f'ADF Statistic: {result[0]:.6f}')
    print(f'p-value: {result[1]:.6f}')
    print('Critical Values:')
    for key, value in result[4].items():
        print(f'\t{key}: {value:.3f}')
    
    if result[1] <= 0.05:
        print("序列是平稳的")
        return True
    else:
        print("序列是非平稳的")
        return False

print("\n原始序列平稳性检验:")
is_stationary = check_stationarity(ts_data['sales'])

# 差分处理
diff_sales = ts_data['sales'].diff().dropna()

plt.subplot(2, 2, 2)
plt.plot(diff_sales.index, diff_sales, marker='o')
plt.title('一阶差分后的时间序列')
plt.xlabel('时间')
plt.ylabel('差分值')
plt.grid(True)

print("\n一阶差分后平稳性检验:")
is_diff_stationary = check_stationarity(diff_sales)

# 绘制ACF和PACF图
plt.subplot(2, 2, 3)
plot_acf(diff_sales, lags=20, ax=plt.gca())
plt.title('自相关函数 (ACF)')

plt.subplot(2, 2, 4)
plot_pacf(diff_sales, lags=20, ax=plt.gca())
plt.title('偏自相关函数 (PACF)')

plt.tight_layout()
plt.show()

# 划分训练集和测试集
train_size = int(0.8 * len(ts_data))
train_data = ts_data['sales'][:train_size]
test_data = ts_data['sales'][train_size:]

print(f"\n训练集大小: {len(train_data)}")
print(f"测试集大小: {len(test_data)}")

# 拟合ARIMA模型
# 这里使用ARIMA(2,1,2)模型，实际应用中需要通过网格搜索确定最优参数
model = ARIMA(train_data, order=(2, 1, 2))
fitted_model = model.fit()

print("\nARIMA模型拟合结果:")
print(fitted_model.summary())

# 模型预测
forecast_steps = len(test_data)
forecast = fitted_model.forecast(steps=forecast_steps)
forecast_ci = fitted_model.get_forecast(steps=forecast_steps).conf_int()

# 计算预测误差
mae = np.mean(np.abs(test_data - forecast))
mse = np.mean((test_data - forecast) ** 2)
rmse = np.sqrt(mse)

print(f"\n预测性能指标:")
print(f"平均绝对误差 (MAE): {mae:.2f}")
print(f"均方误差 (MSE): {mse:.2f}")
print(f"均方根误差 (RMSE): {rmse:.2f}")

# 可视化预测结果
plt.figure(figsize=(14, 8))

# 绘制训练数据、测试数据和预测结果
plt.subplot(2, 1, 1)
plt.plot(train_data.index, train_data, label='训练数据', color='blue')
plt.plot(test_data.index, test_data, label='实际数据', color='green', marker='o')
plt.plot(test_data.index, forecast, label='预测数据', color='red', marker='s')
plt.fill_between(test_data.index, 
                 forecast_ci.iloc[:, 0], 
                 forecast_ci.iloc[:, 1], 
                 color='red', alpha=0.2, label='置信区间')
plt.title('ARIMA时间序列预测结果')
plt.xlabel('时间')
plt.ylabel('销售额')
plt.legend()
plt.grid(True)

# 绘制残差图
residuals = fitted_model.resid
plt.subplot(2, 1, 2)
plt.plot(residuals.index, residuals)
plt.title('模型残差')
plt.xlabel('时间')
plt.ylabel('残差')
plt.grid(True)

plt.tight_layout()
plt.show()

# 残差分析
plt.figure(figsize=(12, 8))

plt.subplot(2, 2, 1)
plt.plot(residuals)
plt.title('残差时间序列')
plt.xlabel('时间')
plt.ylabel('残差')

plt.subplot(2, 2, 2)
plt.hist(residuals, bins=20, alpha=0.7, edgecolor='black')
plt.title('残差分布直方图')
plt.xlabel('残差')
plt.ylabel('频数')

plt.subplot(2, 2, 3)
plot_acf(residuals, lags=20, ax=plt.gca())
plt.title('残差自相关函数')

plt.subplot(2, 2, 4)
from scipy import stats
stats.probplot(residuals, dist="norm", plot=plt)
plt.title('残差Q-Q图')

plt.tight_layout()
plt.show()

# 未来预测
future_steps = 6
future_forecast = fitted_model.forecast(steps=future_steps)
future_ci = fitted_model.get_forecast(steps=future_steps).conf_int()

# 创建未来日期
future_dates = pd.date_range(start=ts_data.index[-1] + pd.DateOffset(months=1), 
                           periods=future_steps, freq='M')

print(f"\n未来{future_steps}个月预测:")
for i, (date, pred, lower, upper) in enumerate(zip(future_dates, future_forecast, 
                                                  future_ci.iloc[:, 0], future_ci.iloc[:, 1])):
    print(f"{date.strftime('%Y-%m')}: {pred:.2f} (置信区间: [{lower:.2f}, {upper:.2f}])")

# 绘制完整预测图
plt.figure(figsize=(14, 6))
plt.plot(ts_data.index, ts_data['sales'], label='历史数据', color='blue')
plt.plot(future_dates, future_forecast, label='未来预测', color='red', marker='o')
plt.fill_between(future_dates, future_ci.iloc[:, 0], future_ci.iloc[:, 1], 
                 color='red', alpha=0.2, label='预测置信区间')
plt.title('ARIMA模型：历史数据与未来预测')
plt.xlabel('时间')
plt.ylabel('销售额')
plt.legend()
plt.grid(True)
plt.show() 